package com.ai.aiagent.advisor;

import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

/**
 * 自定义增强re2拦截器
 */
public class ReReadingAdvisor implements CallAdvisor, StreamAdvisor {

	private final int order;
	public ReReadingAdvisor(int order) {
		this.order = order;
	}

	private ChatClientRequest before(ChatClientRequest advisedRequest) {

		/**
		 * 目前springai alibaba的链式调用不能设置模版prompt
		 * 只能使用spring ai的PromptTemplate来进行结构化模版prompt
		 */
		Map<String, Object> advisedUserParams = new HashMap<>();
		advisedUserParams.put("re2_input_query", advisedRequest.prompt().getUserMessage().getText());

		// 使用prompt模板,重读传入用户的输入信息
		PromptTemplate promptTemplate = new PromptTemplate("""
				{re2_input_query}
				Read the question again: {re2_input_query}
				""");
		// 查看promptTemplate源码可看到不同的将增强后的用户prompt传入chatclient中的方法
		// 查看ChatClientRequest的bulider中prompt，Prompt中初始化的方法

		// 1. Prompt create(Map<String, Object> additionalVariables) 方法可以生成prompt。
		Prompt prompt = promptTemplate.create(advisedUserParams);
		// 2. Message createMessage(Map<String, Object> additionalVariables) 方法可以生成message
		Message message = promptTemplate.createMessage(advisedUserParams);
		// 3. String render(Map<String, Object> additionalVariables) 方法可以生成prompt	内容
		String content = promptTemplate.render(advisedUserParams);
		return ChatClientRequest.builder()
//				.prompt(prompt)
//				.prompt(new Prompt(message))
				.prompt(new Prompt(content))
				.build();


	}

	@Override
	public int getOrder() {
		return this.order;
	}

    @Override
    public String getName() {
		return this.getClass().getSimpleName();
	}

	@Override
	public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
		return callAdvisorChain.nextCall(this.before(chatClientRequest));
	}

	@Override
	public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain streamAdvisorChain) {
		return streamAdvisorChain.nextStream(this.before(chatClientRequest));
	}

	public static Builder builder() {
		return new Builder();
	}

	public static final class Builder {
		private int order = 0;

		private Builder() {
		}

		public Builder order(int order) {
			this.order = order;
			return this;
		}

		public ReReadingAdvisor build() {
			return new ReReadingAdvisor(order);
		}
	}
}